import numpy as np
import torch
def global_dist(array1, array2, type='euclidean'):
    '''Compute the euclidean or cosine distance of all pairs.
    Args:
        array1: numpy array with shape [M1, N]
        array2: numpy array with shape [M2, N]
        type: one of ['cosine','euclidean']
        Return:
            numpy array with shape[M1, M2]
        '''
    assert type in ['cosine', 'euclidean']
    if type == 'consine':
        array1 = normalize(array1, axis=1)
        array2 = normalize(array2, axis=1)
        dist = np.matmul(array1, array2.T)
    else:
        # shape[m1, 1]
        square1 = np.sum(np.square(array1), axis=1)[..., np.newaxis]
        # shape[1, m2]
        square2 = np.sum(np.square(array2), axis=1)[np.newaxis, ...]
        squared_dist = -2 * np.matmul(array1, array2.T) + square1 + square2
        squared_dist[squared_dist<0] = 0
        dist = np.sqrt(squared_dist)
        return dist

def normalize(nparray, order=2, axis=0):
    '''Normalize a N-D numpy array along the specified axis.'''
    norm = np.linalg.norm(nparray, ord=order, axis=axis, keepdims=True)
    return nparray / (norm+np.finfo(np.float32).eps)

def spatial_dist(array1, array2):
    '''Computing the spatial feature reconstruction of all pairs
    Args:
        array1: shape[M1, N, C] M1:the number of query, N:the number of spatial feature, C:the dimension of each spatial feature
        array2：shape[M2, N, C] M2:the number of gallery, N:the number of spatial feature, C: the dimension of each spatial feature.
        :return: numpy array shape[M1 ,M2]
        '''
    kappa = 1e-3
    dist = torch.zeros(len(array1), len(array2))
    dist = dist.cuda()
    for i in range(len(array2)):
        if(i%100 == 0):
            print('{}/{} test batches done'.format(i, len(array2)))
        y = torch.FloatTensor(array2[i])
        y = y.cuda()
        T = kappa * torch.eye(y.size(1))
        T = T.cuda()
        Proj_M = torch.matmul(torch.inverse(torch.matmul(y.t(), y) + T), y.t())  # (Y^{T} * Y + kappa * I)^{-1} * Y^{T}
        for j in range(0, len(array1)):
            temp = array1[j]
            temp = torch.FloatTensor(temp)
            temp = temp.cuda()
            a = torch.matmul(y, torch.matmul(Proj_M, temp)) - temp
            dist[j, i] = torch.pow(a, 2).sum(0).sqrt().mean()
    dist = dist.cpu()
    dist = dist.numpy()
    return dist